import os
from itertools import chain
from typing import Tuple

import cv2
from PIL import Image
from scipy.ndimage import gaussian_filter
from skimage.filters import threshold_otsu
import numpy as np
from tqdm import tqdm

from data.base_dataset import BaseDataset
import abc

from data.preprocess import preprocess_all, generate_hint
from utils.util import mkdirs

IMG_EXTENSIONS = [
    'jpg', 'jpeg',
    'png', 'bmp'
]

isImage = lambda x: x.split('.')[-1].lower() in IMG_EXTENSIONS
listdir = lambda x: [os.path.join(x, f).replace('\\', '/') for f in os.listdir(x)]
isnpz = lambda f: f.split(".")[-1].lower() in ["npz"]

class xDog:
    def __init__(self,
                 γ: float = 0.95,
                 ϕ: float = 1e9,
                 ϵ: float = -1e1,
                 k: float = 4.5,
                 σ: float = 0.3, ) -> None:
        self.γ = γ
        self.ϕ = ϕ
        self.ϵ = ϵ
        self.k = k
        self.σ = σ

    def __call__(self, img) -> np.ndarray:
        x = (img[..., 0] + img[..., 1] + img[..., 2]) / 3

        gaussian_a = gaussian_filter(x, self.σ)
        gaussian_b = gaussian_filter(x, self.σ * self.k)

        dog = gaussian_a - self.γ * gaussian_b

        inf = dog < self.ε
        xdog = inf * 1 + ~inf * (1 - np.tanh(self.φ * dog))

        xdog -= xdog.min()
        xdog /= xdog.max()
        xdog = xdog >= threshold_otsu(xdog)
        xdog = 1 - xdog

        return xdog

# Construct supervised image pairs in the form of hint-colorization
class HintDataset(BaseDataset):
    def __init__(self, opt):
        super(HintDataset, self).__init__(opt)
        self.dataroot = opt.dataroot
        self.skelentonizer_path = opt.skelentonizer_path
        self.sketch_root = os.path.join(self.dataroot,'sketches')
        self.images_root = os.path.join(self.dataroot,'images')
        self.preprocess_root = os.path.join(self.dataroot,'preprocess')
        mkdirs(self.sketch_root)
        mkdirs(self.images_root)
        mkdirs(self.preprocess_root)
        self.image_list, self.sketch_list,self.prep_list = self.build_dataset()
        print(len(self.image_list), len(self.sketch_list), len(self.prep_list))


    def build_dataset(self):
        image_list, sketch_list = self.make_dataset(self.images_root, self.sketch_root)
        preprocess_all(self.images_root, self.preprocess_root, model_path=self.skelentonizer_path)
        prep = chain(*(filter(isnpz, listdir(f)) for f in listdir(self.preprocess_root)))
        prep = sorted(list(prep))
        return image_list, sketch_list, prep

    def __len__(self):
        return min(len(self.sketch_list), len(self.image_list), len(self.prep_list))

    def __getitem__(self, index):
        image_path,sketch_path = self.image_list[index], self.sketch_list[index]
        npz = self.prep_list[index]
        colors, masks = self.load_preprocessed(npz)
        image = self.openImage(image_path)
        sketch = Image.open(sketch_path).convert('L')
        hints = generate_hint(np.array(image)/255, masks)
        hints = Image.fromarray((hints * 255).astype(np.uint8)).convert("RGB")
        post_dict = {
            'image': image,
            'sketch': sketch,
            'hints': hints,
        }
        result = self.postprocess(post_dict)
        return result


    @abc.abstractmethod
    def postprocess(self, input_dict):
        pass

    def openImage(self,root):
        if not os.path.exists(root):
            raise FileNotFoundError('not found file')
        if not isImage(root):
            raise Exception('illegal img file')
        img = Image.open(root).convert("RGB")
        return img

    def load_preprocessed(self, path: str) -> Tuple[np.ndarray, np.ndarray]:
        loader = np.load(path)
        return loader["colors"], loader["masks"]

    def generate_lineart(self, img_path, desc_dir, lineart_generator):
        *_, artist, img = img_path.split("/")
        folder = os.path.join(desc_dir, artist).replace('\\', '/')
        os.makedirs(folder, exist_ok=True)

        sketch_name = f"{img.split('.')[-2]}.png"
        path = os.path.join(folder, sketch_name).replace('\\', '/')

        if os.path.isfile(path):
            return path

        img = self.openImage(img_path)
        img_ndarray = np.array(img) / 255.0
        sketch = (cv2.cvtColor(lineart_generator(img_ndarray).astype(np.float32), cv2.COLOR_GRAY2RGB) * 255).astype(
            np.int32)
        cv2.imwrite(path, sketch)
        return path

    def make_dataset(self, ImageFolder: str, sketchFolder=None):
        if not os.path.isdir(ImageFolder):
            raise Exception("please provide folder of Image")
        images_list = chain(*(filter(isImage, listdir(f)) for f in listdir(ImageFolder)))
        images_list = sorted(list(images_list))

        # if sketchFolder is not None:
        #     sketch_list = chain(*(filter(isimg, listdir(f)) for f in listdir(sketchFolder)))
        #     sketch_list = sorted(list(sketch_list))
        #     return images_list, sketch_list

        if sketchFolder is None:
            sketch_dirname = os.path.abspath(os.path.join(ImageFolder, "..", 'sketches')).replace('\\', '/')
        else:
            sketch_dirname = sketchFolder
        if not os.path.isdir(sketch_dirname):
            os.makedirs(sketch_dirname)
        # sketchs_root = os.path.join(sketch_dirname, folderName).replace('\\', '/')
        pbar = tqdm(images_list)
        line_generator = xDog()
        sketch_list = list()
        for img in pbar:
            path = self.generate_lineart(img, sketch_dirname, line_generator)
            sketch_list.append(path)
        sketch_list = sorted(sketch_list)
        return images_list, sketch_list